#!/usr/bin/env python
# -*- coding: utf-8 -*-

import csv
import os
import random
import sys
import re
from scapy.utils import PcapWriter
from scapy.all import *

class CSV2PCAP:
    def __init__(self) -> None:
        pass

    def build_handshake(self, mac_src, mac_dst, src, dst, sport, dport, pktdump):
        macsrc = mac_src
        macdst = mac_dst
        ipsrc = src
        ipdst = dst
        portsrc = sport
        portdst = dport

        client_isn = random.randint(1024, 10000)
        server_isn = random.randint(1024, 10000)
        syn = Ether(src=macsrc, dst=macdst)/IP(src=ipsrc, dst=ipdst)/TCP(flags="S", sport=portsrc, dport=portdst, seq=client_isn)
        synack = Ether(src=macsrc, dst=macdst)/IP(src=ipdst, dst=ipsrc)/TCP(flags="SA", sport=portdst, dport=portsrc, seq=server_isn, ack=syn.seq+1)
        ack = Ether(src=macsrc, dst=macdst)/IP(src=ipsrc, dst=ipdst)/TCP(flags="A", sport=portsrc, dport=portdst, seq=syn.seq+1, ack=synack.seq+1)
        pktdump.write(syn)
        pktdump.write(synack)
        pktdump.write(ack)
        return (ack.seq, ack.ack)

    def build_finshake(self, mac_src, mac_dst, src, dst, sport, dport, seq, ack, pktdump):
        macsrc = mac_src
        macdst = mac_dst
        ipsrc = src
        ipdst = dst
        portsrc = sport
        portdst = dport
        finAck = Ether(src=macsrc, dst=macdst)/IP(src=ipsrc, dst=ipdst)/TCP(flags="FA", sport=sport, dport=dport, seq=seq, ack=ack)
        finalAck = Ether(src=macsrc, dst=macdst)/IP(src=ipdst, dst=ipsrc)/TCP(flags="A", sport=dport, dport=sport, seq=finAck.ack, ack=finAck.seq+1)
        pktdump.write(finAck)
        pktdump.write(finalAck)

    def chunkstring(self, string, length):
        return (string[0+i:length+i] for i in range(0, len(string), length))

    def make_poop(self, mac_src, mac_dst, src, dst, sport, dport, seq, ack, payload, pktdump):
        segments = []
        if len(payload) > 1460:
            segments = list(self.chunkstring(payload, 1460))
        else:
            segments.append(payload)

        macsrc = mac_src
        macdst = mac_dst
        ipsrc = src
        ipdst = dst
        portsrc = sport
        portdst = dport
        for segment in segments:
            p = Ether(src=macsrc, dst=macdst)/IP(src=ipsrc, dst=ipdst)/TCP(flags="PA", sport=sport, dport=dport, seq=seq, ack=ack)/segment
            returnAck = Ether(src=macsrc, dst=macdst)/IP(src=ipdst, dst=ipsrc)/TCP(flags="A", sport=dport, dport=sport, seq=p.ack, ack=(p.seq + len(p[Raw])))
            seq = returnAck.ack
            ack = returnAck.seq
            pktdump.write(p)
            pktdump.write(returnAck)
        return (returnAck.seq, returnAck.ack)

    def convert_headers(self, headers_str):
        """将CSV中的请求头或响应头字符串转换为HTTP协议标准格式"""
        if not headers_str or headers_str.strip() == "":
            return ""
        
        # 首先规范化空白字符
        headers_str = headers_str.strip()
        headers_str = re.sub(r'\s+', ' ', headers_str)
        
        # 常见的HTTP请求方法
        http_methods = ['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH', 'CONNECT', 'TRACE']
        
        # 尝试提取请求行或状态行
        request_line = ""
        result_headers = []
        
        # 检查是否以HTTP方法开头
        for method in http_methods:
            if headers_str.startswith(method + ' '):
                parts = headers_str.split(' ', 3)
                if len(parts) >= 3 and 'HTTP/' in parts[2]:
                    request_line = f"{parts[0]} {parts[1]} {parts[2]}"
                    # 移除请求行
                    headers_str = headers_str[len(request_line):].strip()
                    result_headers.append(request_line)
                    break
        
        # 如果以HTTP/开头，可能是响应状态行
        if not request_line and headers_str.startswith('HTTP/'):
            parts = headers_str.split(' ', 3)
            if len(parts) >= 3:
                request_line = f"{parts[0]} {parts[1]} {parts[2]}"
                # 移除状态行
                headers_str = headers_str[len(request_line):].strip()
                result_headers.append(request_line)
        
        # 使用更可靠的方法：直接查找所有已知的HTTP头字段并分割
        # 常见的HTTP头字段名
        common_headers = [
            'Host:', 'User-Agent:', 'Accept:', 'Accept-Encoding:', 'Accept-Language:', 
            'Content-Type:', 'Content-Length:', 'Connection:', 'Cookie:', 'Referer:',
            'Authorization:', 'Date:', 'Server:', 'Set-Cookie:', 'X-Forwarded-For:',
            'X-Real-IP:', 'WZWS-RAY:' 
        ]
        
        # 使用直接的替换方法
        if headers_str:
            # 首先检查第一个头部
            found_first_header = False
            for header in common_headers:
                if headers_str.startswith(header):
                    found_first_header = True
                    break
            
            if not found_first_header:
                # 如果第一个部分不是标准头部，可能需要特殊处理
                for header in common_headers:
                    pos = headers_str.find(' ' + header)
                    if pos >= 0:
                        # 将第一个头部前的所有内容视为一个独立的头部
                        first_part = headers_str[:pos].strip()
                        if first_part:
                            result_headers.append(first_part)
                        headers_str = headers_str[pos+1:] # +1 跳过空格
                        break
            
            # 为剩余的头部添加换行符
            processed = headers_str
            for header in common_headers:
                # 只替换前面有空格的头部字段（防止替换字段值中的相似字符串）
                processed = re.sub(f' ({re.escape(header)})', r'\r\n\1', processed)
            
            # 添加最后处理的部分
            if processed:
                if result_headers:
                    # 如果已经有其他头部，将处理后的内容附加上去
                    result_headers.append(processed)
                else:
                    # 如果这是唯一的内容，直接返回
                    return processed
        
        # 返回处理后的HTTP头
        return '\r\n'.join(result_headers)

    def process_csv(self, csv_file, output_dir="pcap_output"):
        """处理CSV文件，为每行生成一个PCAP文件"""
        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"已创建输出目录: {output_dir}")
        
        with open(csv_file, 'r', encoding='utf-8', errors='ignore') as f:
            reader = csv.reader(f)
            # 跳过表头
            header = next(reader)
            
            # 确定列索引
            src_ip_idx = header.index('源IPV4') if '源IPV4' in header else 2
            src_port_idx = header.index('源端口') if '源端口' in header else 4
            dst_ip_idx = header.index('目的IPV4') if '目的IPV4' in header else 7
            dst_port_idx = header.index('目的端口') if '目的端口' in header else 9
            req_header_idx = header.index('HTTP请求头') if 'HTTP请求头' in header else 31
            req_body_idx = header.index('HTTP请求体') if 'HTTP请求体' in header else 34
            resp_header_idx = header.index('HTTP响应头') if 'HTTP响应头' in header else 36
            resp_body_idx = header.index('HTTP响应体') if 'HTTP响应体' in header else 39
            
            for idx, row in enumerate(reader, 1):
                try:
                    # 跳过空行
                    if not row or len(row) < max(src_ip_idx, src_port_idx, dst_ip_idx, dst_port_idx, 
                                              req_header_idx, req_body_idx, resp_header_idx, resp_body_idx) + 1:
                        print(f"跳过行 {idx}: 数据不完整")
                        continue
                    
                    # 解析CSV行数据
                    src_ip = row[src_ip_idx].strip()
                    try:
                        src_port = int(row[src_port_idx].strip())
                    except ValueError:
                        print(f"行 {idx} 的源端口无效，使用默认值80")
                        src_port = 80
                        
                    dst_ip = row[dst_ip_idx].strip()
                    try:
                        dst_port = int(row[dst_port_idx].strip())
                    except ValueError:
                        print(f"行 {idx} 的目标端口无效，使用默认值80")
                        dst_port = 80
                    
                    # 检查IP地址是否有效
                    if not self.validate_ip(src_ip) or not self.validate_ip(dst_ip):
                        print(f"跳过行 {idx}: IP地址无效")
                        continue
                    
                    # 提取并转换HTTP头和正文
                    http_req_header = self.convert_headers(row[req_header_idx])
                    http_req_body = row[req_body_idx].strip() if row[req_body_idx] else ""
                    http_resp_header = self.convert_headers(row[resp_header_idx])
                    http_resp_body = row[resp_body_idx].strip() if row[resp_body_idx] else ""
                    
                    # 构建HTTP请求
                    # 1. 处理请求头，确保以\r\n结尾
                    req_payload = http_req_header.rstrip('\r\n')
                    # 2. 添加\r\n\r\n分隔符（无论请求体是否为空）
                    req_payload += '\r\n\r\n'
                    # 3. 添加请求体
                    req_payload += http_req_body
                    # 判断http_req_body不为空的话，添加\r\n\r\n分隔符
                    if http_req_body:
                        req_payload += '\r\n\r\n'
                    
                    # 构建HTTP响应
                    # 1. 处理响应头，确保以\r\n结尾
                    resp_payload = http_resp_header.rstrip('\r\n')
                    # 2. 添加\r\n\r\n分隔符（无论响应体是否为空）
                    resp_payload += '\r\n\r\n'
                    # 3. 添加响应体
                    resp_payload += http_resp_body
                    # 判断http_resp_body不为空的话，添加\r\n\r\n分隔符
                    if http_resp_body:
                        resp_payload += '\r\n\r\n'
                    
                    # 生成PCAP文件名
                    pcap_filename = os.path.join(output_dir, f"{idx}_{src_ip}_{src_port}.pcap")
                    
                    # 创建PCAP文件
                    mac_src = row[1].strip() if len(row) > 1 and row[1].strip() else "00:0c:29:65:b0:ba"
                    mac_dst = row[6].strip() if len(row) > 6 and row[6].strip() else "00:50:56:c0:00:01"
                    
                    self.create_pcap(pcap_filename, src_ip, dst_ip, src_port, dst_port, req_payload, resp_payload, mac_src, mac_dst)
                    
                    print(f"已生成PCAP文件: {pcap_filename}")
                    
                except Exception as e:
                    print(f"处理行 {idx} 时出错: {e}")
    
    def validate_ip(self, ip):
        """验证IP地址是否有效"""
        if not ip or ip == "":
            return False
        try:
            parts = ip.split('.')
            if len(parts) != 4:
                return False
            for part in parts:
                if not part.isdigit() or not 0 <= int(part) <= 255:
                    return False
            return True
        except:
            return False
    
    def create_pcap(self, pcap_filename, src_ip, dst_ip, src_port, dst_port, req_payload, resp_payload, mac_src="00:0c:29:65:b0:ba", mac_dst="00:50:56:c0:00:01"):
        """创建一个PCAP文件，包含完整的TCP/HTTP会话"""
        # 创建PCAP文件写入器
        pktdump = PcapWriter(pcap_filename, sync=True)
        
        # TCP三次握手
        (seq, ack) = self.build_handshake(mac_src, mac_dst, src_ip, dst_ip, src_port, dst_port, pktdump)
        
        # 发送HTTP请求并接收ACK
        (seq, ack) = self.make_poop(mac_src, mac_dst, src_ip, dst_ip, src_port, dst_port, seq, ack, req_payload, pktdump)
        
        # 发送HTTP响应并接收ACK
        (seq, ack) = self.make_poop(mac_src, mac_dst, dst_ip, src_ip, dst_port, src_port, seq, ack, resp_payload, pktdump)
        
        # TCP四次挥手
        self.build_finshake(mac_src, mac_dst, src_ip, dst_ip, src_port, dst_port, seq, ack, pktdump)
        
        # 关闭PCAP文件
        pktdump.close()

if __name__ == "__main__":
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print("用法: python csv2pcap.py <csv文件路径> [输出目录]")
        sys.exit(1)
    
    csv_file = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) == 3 else "pcap_output"
    
    if not os.path.exists(csv_file):
        print(f"错误: 文件 {csv_file} 不存在")
        sys.exit(1)
    
    processor = CSV2PCAP()
    processor.process_csv(csv_file, output_dir)
    print(f"所有PCAP文件已生成完成到目录: {output_dir}!") 